{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "163e1754-2bda-47b1-b40f-d07e03450788",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pytblis\n",
    "import torch\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1fa45174-7885-44c4-8fb6-f894a80dbee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "nfull = 128\n",
    "nao = 96\n",
    "e_full = np.cos(np.arange(nfull**4) + 0.2).reshape(nfull, nfull, nfull, nfull)\n",
    "e = e_full[:nao, :nao, :nao, :nao]\n",
    "e_torch = torch.asarray(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "68c78b52-a386-4ef0-a17e-0521fc49bb32",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fp(arr):\n",
    "    return np.cos(np.arange(arr.size)) @ arr.reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "042162b6-5986-4e00-82ff-7a0d701dc411",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fp_torch(arr):\n",
    "    return torch.cos(torch.arange(np.prod(list(arr.size())), dtype=torch.double)) @ arr.reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "980fca0e-72a0-4241-957c-8c55dbf2e186",
   "metadata": {},
   "outputs": [],
   "source": [
    "subscripts_list = [\n",
    "    \"abxy, xycd -> abcd\",  # naive gemm case, 2 * n^6\n",
    "    \"axyz, xyzb -> ab\",    # naive gemm case, 2 * n^5\n",
    "    \"axyz, bxyz -> ab\",    # naive syrk case,     n^5\n",
    "    \"axyz, ybzx -> ab\",    # comp  gemm case, 2 * n^5\n",
    "    \"axby, yacx -> abc\",   # batch gemm case, 2 * n^5\n",
    "    \"xpay, aybx -> ab\",    # complicate case, 2 * n^4\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "554f3ebc-6f31-4388-b4d5-bcf226ce4f8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NumPy einsum\n",
      "Subscripts: abxy, xycd -> abcd\n",
      "elapsed time:     7.306426 sec (avg of  5 repeats)\n",
      "fingerprint :     -20.188290390819\n",
      "Subscripts: axyz, xyzb -> ab\n",
      "elapsed time:     0.776790 sec (avg of 20 repeats)\n",
      "fingerprint :      20.343405116707\n",
      "Subscripts: axyz, bxyz -> ab\n",
      "elapsed time:     0.455486 sec (avg of 20 repeats)\n",
      "fingerprint : -200211.721311474335\n",
      "Subscripts: axyz, ybzx -> ab\n",
      "elapsed time:     0.725011 sec (avg of 20 repeats)\n",
      "fingerprint :       0.274707781823\n",
      "Subscripts: axby, yacx -> abc\n",
      "elapsed time:    27.076382 sec (avg of  1 repeats)\n",
      "fingerprint :       0.466623082298\n",
      "Subscripts: xpay, aybx -> ab\n",
      "elapsed time:   248.522274 sec (avg of  1 repeats)\n",
      "fingerprint :       0.134542958876\n"
     ]
    }
   ],
   "source": [
    "print(\"NumPy einsum\")\n",
    "repeat_list = [5, 20, 20, 20, 1, 1]\n",
    "for subscripts, nrepeat in zip(subscripts_list, repeat_list):\n",
    "    print(f\"Subscripts: {subscripts}\")\n",
    "    t = time.time()\n",
    "    for _ in range(nrepeat):\n",
    "        v = np.einsum(subscripts, e, e, optimize=True)\n",
    "    print(f\"elapsed time: {(time.time() - t) / nrepeat:12.6f} sec (avg of {nrepeat:2d} repeats)\")\n",
    "    print(f\"fingerprint : {fp(v):20.12f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b408c155-71c5-4b8d-9a8b-40911b2d1209",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTBLIS einsum\n",
      "Subscripts: abxy, xycd -> abcd\n",
      "elapsed time:     1.951329 sec (avg of  5 repeats)\n",
      "fingerprint :     -20.188290390824\n",
      "Subscripts: axyz, xyzb -> ab\n",
      "elapsed time:     0.141577 sec (avg of 20 repeats)\n",
      "fingerprint :      20.343405116705\n",
      "Subscripts: axyz, bxyz -> ab\n",
      "elapsed time:     0.114764 sec (avg of 20 repeats)\n",
      "fingerprint : -200211.721311338129\n",
      "Subscripts: axyz, ybzx -> ab\n",
      "elapsed time:     0.139257 sec (avg of 20 repeats)\n",
      "fingerprint :       0.274707781825\n",
      "Subscripts: axby, yacx -> abc\n",
      "elapsed time:    27.224720 sec (avg of  1 repeats)\n",
      "fingerprint :       0.466623082298\n",
      "Subscripts: xpay, aybx -> ab\n",
      "elapsed time:   249.175256 sec (avg of  1 repeats)\n",
      "fingerprint :       0.134542958876\n"
     ]
    }
   ],
   "source": [
    "print(\"PyTBLIS einsum\")\n",
    "repeat_list = [5, 20, 20, 20, 1, 1]\n",
    "for subscripts, nrepeat in zip(subscripts_list, repeat_list):\n",
    "    print(f\"Subscripts: {subscripts}\")\n",
    "    t = time.time()\n",
    "    for _ in range(nrepeat):\n",
    "        v = pytblis.einsum(subscripts, e, e, optimize=\"greedy\")\n",
    "    print(f\"elapsed time: {(time.time() - t) / nrepeat:12.6f} sec (avg of {nrepeat:2d} repeats)\")\n",
    "    print(f\"fingerprint : {fp(v):20.12f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e27aeb8f-5f3f-48eb-aabf-e9c9ba9eaa50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch einsum\n",
      "Subscripts: abxy, xycd -> abcd\n",
      "elapsed time:     2.105942 sec (avg of  5 repeats)\n",
      "fingerprint :     -20.188290390819\n",
      "Subscripts: axyz, xyzb -> ab\n",
      "elapsed time:     0.204239 sec (avg of 20 repeats)\n",
      "fingerprint :      20.343405116707\n",
      "Subscripts: axyz, bxyz -> ab\n",
      "elapsed time:     0.211401 sec (avg of 20 repeats)\n",
      "fingerprint : -200211.721311473870\n",
      "Subscripts: axyz, ybzx -> ab\n",
      "elapsed time:     0.406712 sec (avg of 20 repeats)\n",
      "fingerprint :       0.274707781823\n",
      "Subscripts: axby, yacx -> abc\n",
      "elapsed time:     0.263642 sec (avg of 20 repeats)\n",
      "fingerprint :       0.466623082310\n",
      "Subscripts: xpay, aybx -> ab\n",
      "elapsed time:     0.147302 sec (avg of 20 repeats)\n",
      "fingerprint :       0.134542958873\n"
     ]
    }
   ],
   "source": [
    "print(\"PyTorch einsum\")\n",
    "repeat_list = [5, 20, 20, 20, 20, 20]\n",
    "for subscripts, nrepeat in zip(subscripts_list, repeat_list):\n",
    "    print(f\"Subscripts: {subscripts}\")\n",
    "    t = time.time()\n",
    "    for _ in range(nrepeat):\n",
    "        v = torch.einsum(subscripts, e_torch, e_torch)\n",
    "    print(f\"elapsed time: {(time.time() - t) / nrepeat:12.6f} sec (avg of {nrepeat:2d} repeats)\")\n",
    "    print(f\"fingerprint : {fp_torch(v):20.12f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
